import cv2
import mediapipe as mp
from utils.data_aug import position_project
from utils.get_distance import cal_rectangle
from utils.train_val_split import ratio_split
from utils.dir_operation import clear_folder,change_folder_permissions,zip_folder
from pathlib import Path
import numpy as np
import os

def start_record(root, category, video_path):
    Path(root).mkdir(parents=True, exist_ok=True)
    p = Path(root) / category
    p.mkdir(parents=True, exist_ok=True)
    # 创建 Path 对象
    path_obj = Path(video_path)
    # 获取去掉扩展名的文件名
    name = path_obj.stem
    cap = cv2.VideoCapture(video_path)
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        # 水平翻转图像
        frame = cv2.flip(frame, 1)
        # 将图像从 BGR 转换为 RGB
        image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        image.flags.writeable = False
        # 处理图像以检测手部关键点
        results = hands.process(image)
        # 将图像从 RGB 转换回 BGR
        image.flags.writeable = True
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        h, w, _ = image.shape
        # 绘制手部关键点和连接线
        if results.multi_hand_landmarks:
            for hand_landmarks in results.multi_hand_landmarks:
                min_x, min_y, max_x, max_y = cal_rectangle(hand_landmarks)
                # 将归一化的坐标转换为像素坐标
                min_x, min_y = int(min_x * w), int(min_y * h)
                max_x, max_y  = int(max_x * w), int(max_y * h)
                # xy_dis = get_all_distances(hand_landmarks)
                npy_path = p / (name)
                distance, indices = position_project(min_x, min_y, max_x, max_y, hand_landmarks)
                for i in range(len(indices)):
                    abs_path = str(npy_path)  + "_" + indices[i] + ".npy"
                    np.save(abs_path, distance[i])
                cv2.rectangle(image, (min_x, min_y), (max_x, max_y), (0, 255, 0), 2)
                mp_drawing.draw_landmarks(
                    image, hand_landmarks, mp_hands.HAND_CONNECTIONS)
        
        # 调整图像大小
        resized_image = cv2.resize(image, (display_width, display_height))
        # 显示图像
        cv2.imshow('Hand Landmark Detection', resized_image)
        if cv2.waitKey(5) & 0xFF == 27:  # 按下 ESC 键退出
            break
    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    data_folder = 'data'
    points_folder = 'points'
    clear_folder(data_folder)
    clear_folder(points_folder)
    mp_hands = mp.solutions.hands
    hands = mp_hands.Hands(static_image_mode=False, max_num_hands=2, min_detection_confidence=0.5)
    mp_drawing = mp.solutions.drawing_utils
    cap = cv2.VideoCapture(0)
    # 设置显示窗口的宽度和高度
    display_width = 1000
    display_height = 800
    # 创建一个可调整大小的窗口
    cv2.namedWindow('Hand Landmark Detection', cv2.WINDOW_NORMAL)
    cv2.resizeWindow('Hand Landmark Detection', display_width, display_height)
    video_path = os.path.join(os.getcwd(), 'video')
    for hand_file in os.listdir(video_path):
        start_record(data_folder, hand_file.split("_")[0], os.path.join(video_path, hand_file))
    ratio_split(data_folder, points_folder, 0.8)
    change_folder_permissions(data_folder, 0o777)
    change_folder_permissions(points_folder, 0o777)
    zip_folder(points_folder, points_folder)
    change_folder_permissions(points_folder+".zip", 0o777)